#include "bzbfile/elfparser.h"

#include "bzbfile/file.h"
#include "bzbutil/platform.h"
#include "bzbstring.h"

#include <fstream>
#include <exception>


void load_name_str(std::vector<char> data, std::string &name, uint32_t offset) {
	name = "";
	size_t i = offset;
	while(data[i] != '\0' && i < data.size()) {
		name += data[i];
		i++;
	}
}


uint8_t bzbfile::get_elf_class(const std::string &filename) {

	uint64_t filesize = bzbfile::size(filename);

	// 52 is 32-bit ELF's header size
	if(filesize < 52) {
		throw std::runtime_error("ERROR Unable to read ELF EI_class of " + filename +
								 " because the file is corrupted (file size < header size)");
	}

	std::ifstream stream(filename, std::ios::binary);
	char ei_class;

	if(stream) {

		stream.seekg(ELF_EI_CLASS);
		stream.read(&ei_class, 1);
		stream.close();

		return (uint8_t) ei_class;
	} else {
		bzbutil::throw_sys_error(errno);
	}

	return 0;
}


bzbfile::ELF32 bzbfile::parse_elf32(const std::string &filename) {

	uint64_t filesize = bzbfile::size(filename);


	// Parse header

	if(filesize < 52) {
		throw std::runtime_error("ERROR Unable to parse \'" + filename +
								"\' because the file is corrupted (file size < header size)");
	}

	std::vector<char> header = bzbfile::read_bin(filename, 52);
	ELF32 res;
	res.e_ident[ELF_EI_MAG0] = header[ELF_EI_MAG0];
	res.e_ident[ELF_EI_MAG1] = header[ELF_EI_MAG1];
	res.e_ident[ELF_EI_MAG2] = header[ELF_EI_MAG2];
	res.e_ident[ELF_EI_MAG3] = header[ELF_EI_MAG3];
	res.e_ident[ELF_EI_CLASS] = header[ELF_EI_CLASS];
	res.e_ident[ELF_EI_DATA] = header[ELF_EI_DATA];
	res.e_ident[ELF_EI_VERSION] = header[ELF_EI_VERSION];
	res.e_ident[ELF_EI_OSABI] = header[ELF_EI_OSABI];
	res.e_ident[ELF_EI_ABIVERSION] = header[ELF_EI_ABIVERSION];

	for(size_t i = 0; i < 7; i++) {
		res.e_ident[ELF_EI_PAD + i] = header[ELF_EI_PAD + i];
	}

	if(res.e_ident[ELF_EI_DATA] == ELF_LITTLE_ENDIAN) {

		readint(header, 0x10, &res.e_type);
		readint(header, 0x12, &res.e_machine);
		readint(header, 0x14, &res.e_version);
		readint(header, 0x18, &res.e_entry);
		readint(header, 0x1C, &res.e_phoff);
		readint(header, 0x20, &res.e_shoff);
		readint(header, 0x24, &res.e_flags);
		readint(header, 0x28, &res.e_ehsize);
		readint(header, 0x2A, &res.e_phentsize);
		readint(header, 0x2C, &res.e_phnum);
		readint(header, 0x2E, &res.e_shentsize);
		readint(header, 0x30, &res.e_shnum);
		readint(header, 0x32, &res.e_shstrndx);

	} else if(res.e_ident[ELF_EI_DATA] == ELF_BIG_ENDIAN) {

		readint_be(header, 0x10, &res.e_type);
		readint_be(header, 0x12, &res.e_machine);
		readint_be(header, 0x14, &res.e_version);
		readint_be(header, 0x18, &res.e_entry);
		readint_be(header, 0x1C, &res.e_phoff);
		readint_be(header, 0x20, &res.e_shoff);
		readint_be(header, 0x24, &res.e_flags);
		readint_be(header, 0x28, &res.e_ehsize);
		readint_be(header, 0x2A, &res.e_phentsize);
		readint_be(header, 0x2C, &res.e_phnum);
		readint_be(header, 0x2E, &res.e_shentsize);
		readint_be(header, 0x30, &res.e_shnum);
		readint_be(header, 0x32, &res.e_shstrndx);
	}

	// Parse program headers

	if(res.e_phnum && res.e_phoff) {

		if(filesize < (res.e_phnum * res.e_phentsize + res.e_phoff)) {

			header.clear();

			throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted " +
									"(file size < (e_phnum * e_phentsize + e_phoff))");
		}

		std::vector<char> progheaders = bzbfile::read_bin(filename,
														res.e_phnum * res.e_phentsize,
														res.e_phoff);

		if(progheaders.size() < res.e_phnum * res.e_phentsize) {

			header.clear();
			progheaders.clear();

			throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted " +
									"(program headers are corrupted)");
		}

		res.prog_headers.resize(res.e_phnum);

		for(size_t i = 0; i < res.prog_headers.size(); i++) {

			if(res.e_ident[ELF_EI_DATA] == ELF_LITTLE_ENDIAN) {

				readint(progheaders, 0x00 + (i * res.e_phentsize), &res.prog_headers[i].p_type);
				readint(progheaders, 0x04 + (i * res.e_phentsize), &res.prog_headers[i].p_offset);
				readint(progheaders, 0x08 + (i * res.e_phentsize), &res.prog_headers[i].p_vaddr);
				readint(progheaders, 0x0C + (i * res.e_phentsize), &res.prog_headers[i].p_paddr);
				readint(progheaders, 0x10 + (i * res.e_phentsize), &res.prog_headers[i].p_filesz);
				readint(progheaders, 0x14 + (i * res.e_phentsize), &res.prog_headers[i].p_memsz);
				readint(progheaders, 0x18 + (i * res.e_phentsize), &res.prog_headers[i].p_flags);
				readint(progheaders, 0x1C + (i * res.e_phentsize), &res.prog_headers[i].p_align);

			} else {

				readint_be(progheaders, 0x00 + (i * res.e_phentsize), &res.prog_headers[i].p_type);
				readint_be(progheaders, 0x04 + (i * res.e_phentsize), &res.prog_headers[i].p_offset);
				readint_be(progheaders, 0x08 + (i * res.e_phentsize), &res.prog_headers[i].p_vaddr);
				readint_be(progheaders, 0x0C + (i * res.e_phentsize), &res.prog_headers[i].p_paddr);
				readint_be(progheaders, 0x10 + (i * res.e_phentsize), &res.prog_headers[i].p_filesz);
				readint_be(progheaders, 0x14 + (i * res.e_phentsize), &res.prog_headers[i].p_memsz);
				readint_be(progheaders, 0x18 + (i * res.e_phentsize), &res.prog_headers[i].p_flags);
				readint_be(progheaders, 0x1C + (i * res.e_phentsize), &res.prog_headers[i].p_align);

			}
		}
	}

	if(res.e_shoff && res.e_shnum) {

		if(filesize < (res.e_shnum * res.e_shentsize + res.e_shoff)) {

			header.clear();

			throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted " +
									"(file size < (e_shnum * e_shentsize + e_shoff))");
		}

		std::vector<char> secheaders = bzbfile::read_bin(filename,
														res.e_shnum * res.e_shentsize,
														res.e_shoff);

		if(secheaders.size() < res.e_shnum * res.e_shentsize) {

			header.clear();
			secheaders.clear();

			throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted " +
									"(section headers are corrupted)");
		}

		res.sec_headers.resize(res.e_shnum);

		for(size_t i = 0; i < res.sec_headers.size(); i++) {

			if(res.e_ident[ELF_EI_DATA] == ELF_LITTLE_ENDIAN) {

				readint(secheaders, 0x00 + (i * res.e_shentsize), &res.sec_headers[i].sh_name);
				readint(secheaders, 0x04 + (i * res.e_shentsize), &res.sec_headers[i].sh_type);
				readint(secheaders, 0x08 + (i * res.e_shentsize), &res.sec_headers[i].sh_flags);
				readint(secheaders, 0x0C + (i * res.e_shentsize), &res.sec_headers[i].sh_addr);
				readint(secheaders, 0x10 + (i * res.e_shentsize), &res.sec_headers[i].sh_offset);
				readint(secheaders, 0x14 + (i * res.e_shentsize), &res.sec_headers[i].sh_size);
				readint(secheaders, 0x18 + (i * res.e_shentsize), &res.sec_headers[i].sh_link);
				readint(secheaders, 0x1C + (i * res.e_shentsize), &res.sec_headers[i].sh_info);
				readint(secheaders, 0x20 + (i * res.e_shentsize), &res.sec_headers[i].sh_addralign);
				readint(secheaders, 0x28 + (i * res.e_shentsize), &res.sec_headers[i].sh_entsize);

			} else {

				readint_be(secheaders, 0x00 + (i * res.e_shentsize), &res.sec_headers[i].sh_name);
				readint_be(secheaders, 0x04 + (i * res.e_shentsize), &res.sec_headers[i].sh_type);
				readint_be(secheaders, 0x08 + (i * res.e_shentsize), &res.sec_headers[i].sh_flags);
				readint_be(secheaders, 0x0C + (i * res.e_shentsize), &res.sec_headers[i].sh_addr);
				readint_be(secheaders, 0x10 + (i * res.e_shentsize), &res.sec_headers[i].sh_offset);
				readint_be(secheaders, 0x14 + (i * res.e_shentsize), &res.sec_headers[i].sh_size);
				readint_be(secheaders, 0x18 + (i * res.e_shentsize), &res.sec_headers[i].sh_link);
				readint_be(secheaders, 0x1C + (i * res.e_shentsize), &res.sec_headers[i].sh_info);
				readint_be(secheaders, 0x20 + (i * res.e_shentsize), &res.sec_headers[i].sh_addralign);
				readint_be(secheaders, 0x28 + (i * res.e_shentsize), &res.sec_headers[i].sh_entsize);

			}
		}
	}

	if(res.e_shstrndx &&
	   res.e_shnum >= res.e_shstrndx) {

		if(res.sec_headers[res.e_shstrndx].sh_type != ELF_SHT_STRTAB) {

			header.clear();

			throw std::runtime_error("ERROR Unable to parse \'" + filename +
						"\' because the file is corrupted (e_shstrndx points to invalid section)");
		}

		if(filesize < res.sec_headers[res.e_shstrndx].sh_offset +
					  res.sec_headers[res.e_shstrndx].sh_size) {

			header.clear();
			
			throw std::runtime_error("ERROR Unable to parse \'" + filename +
						"\' because the file is corrupted (.shstrtab section header has invalid data)");
		}

		std::vector<char> shstrtab = bzbfile::read_bin(filename,
														res.sec_headers[res.e_shstrndx].sh_size,
														res.sec_headers[res.e_shstrndx].sh_offset);

		for (size_t i = 0; i < res.sec_headers.size(); ++i) {

			header.clear();
			shstrtab.clear();

			if(res.sec_headers[i].sh_name > res.sec_headers[res.e_shstrndx].sh_size) {
				throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted (invalid sh_name value)");
			}

			load_name_str(shstrtab, res.sec_headers[i].name, res.sec_headers[i].sh_name);
		}

	}

	return res;
}

bzbfile::ELF64 bzbfile::parse_elf64(const std::string &filename) {

	uint64_t filesize = bzbfile::size(filename);

	if(filesize < 64) {
		throw std::runtime_error("ERROR Unable to parse \'" + filename +
								"\' because the file is corrupted (file size < header size)");
	}

	std::vector<char> header = bzbfile::read_bin(filename, 64);
	ELF64 res;
	res.e_ident[ELF_EI_MAG0] = header[ELF_EI_MAG0];
	res.e_ident[ELF_EI_MAG1] = header[ELF_EI_MAG1];
	res.e_ident[ELF_EI_MAG2] = header[ELF_EI_MAG2];
	res.e_ident[ELF_EI_MAG3] = header[ELF_EI_MAG3];
	res.e_ident[ELF_EI_CLASS] = header[ELF_EI_CLASS];
	res.e_ident[ELF_EI_DATA] = header[ELF_EI_DATA];
	res.e_ident[ELF_EI_VERSION] = header[ELF_EI_VERSION];
	res.e_ident[ELF_EI_OSABI] = header[ELF_EI_OSABI];
	res.e_ident[ELF_EI_ABIVERSION] = header[ELF_EI_ABIVERSION];

	for(size_t i = 0; i < 7; i++) {
		res.e_ident[ELF_EI_PAD + i] = header[ELF_EI_PAD + i];
	}

	if(res.e_ident[ELF_EI_DATA] == ELF_LITTLE_ENDIAN) {

		readint(header, 0x10, &res.e_type);
		readint(header, 0x12, &res.e_machine);
		readint(header, 0x14, &res.e_version);
		readint(header, 0x18, &res.e_entry);
		readint(header, 0x20, &res.e_phoff);
		readint(header, 0x28, &res.e_shoff);
		readint(header, 0x30, &res.e_flags);
		readint(header, 0x34, &res.e_ehsize);
		readint(header, 0x36, &res.e_phentsize);
		readint(header, 0x38, &res.e_phnum);
		readint(header, 0x3A, &res.e_shentsize);
		readint(header, 0x3C, &res.e_shnum);
		readint(header, 0x3E, &res.e_shstrndx);

	} else if(res.e_ident[ELF_EI_DATA] == ELF_BIG_ENDIAN) {

		readint_be(header, 0x10, &res.e_type);
		readint_be(header, 0x12, &res.e_machine);
		readint_be(header, 0x14, &res.e_version);
		readint_be(header, 0x18, &res.e_entry);
		readint_be(header, 0x20, &res.e_phoff);
		readint_be(header, 0x28, &res.e_shoff);
		readint_be(header, 0x30, &res.e_flags);
		readint_be(header, 0x34, &res.e_ehsize);
		readint_be(header, 0x36, &res.e_phentsize);
		readint_be(header, 0x38, &res.e_phnum);
		readint_be(header, 0x3A, &res.e_shentsize);
		readint_be(header, 0x3C, &res.e_shnum);
		readint_be(header, 0x3E, &res.e_shstrndx);
	}


	// Parse program headers

	if(res.e_phoff && res.e_phnum) {

		if(filesize < (res.e_phnum * res.e_phentsize + res.e_phoff)) {
			throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted " +
									"(file size < (e_phnum * e_phentsize + e_phoff))");
		}

		std::vector<char> progheaders = bzbfile::read_bin(filename,
														res.e_phnum * res.e_phentsize,
														res.e_phoff);

		if(progheaders.size() < res.e_phnum * res.e_phentsize) {
			throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted " +
									"(program headers are corrupted)");
		}

		res.prog_headers.resize(res.e_phnum);

		for(size_t i = 0; i < res.prog_headers.size(); i++) {

			if(res.e_ident[ELF_EI_DATA] == ELF_LITTLE_ENDIAN) {

				readint(progheaders, 0x00 + (i * res.e_phentsize), &res.prog_headers[i].p_type);
				readint(progheaders, 0x04 + (i * res.e_phentsize), &res.prog_headers[i].p_flags);
				readint(progheaders, 0x08 + (i * res.e_phentsize), &res.prog_headers[i].p_offset);
				readint(progheaders, 0x10 + (i * res.e_phentsize), &res.prog_headers[i].p_vaddr);
				readint(progheaders, 0x18 + (i * res.e_phentsize), &res.prog_headers[i].p_paddr);
				readint(progheaders, 0x20 + (i * res.e_phentsize), &res.prog_headers[i].p_filesz);
				readint(progheaders, 0x28 + (i * res.e_phentsize), &res.prog_headers[i].p_memsz);
				readint(progheaders, 0x30 + (i * res.e_phentsize), &res.prog_headers[i].p_align);

			} else {

				readint_be(progheaders, 0x00 + (i * res.e_phentsize), &res.prog_headers[i].p_type);
				readint_be(progheaders, 0x04 + (i * res.e_phentsize), &res.prog_headers[i].p_flags);
				readint_be(progheaders, 0x08 + (i * res.e_phentsize), &res.prog_headers[i].p_offset);
				readint_be(progheaders, 0x10 + (i * res.e_phentsize), &res.prog_headers[i].p_vaddr);
				readint_be(progheaders, 0x18 + (i * res.e_phentsize), &res.prog_headers[i].p_paddr);
				readint_be(progheaders, 0x20 + (i * res.e_phentsize), &res.prog_headers[i].p_filesz);
				readint_be(progheaders, 0x28 + (i * res.e_phentsize), &res.prog_headers[i].p_memsz);
				readint_be(progheaders, 0x30 + (i * res.e_phentsize), &res.prog_headers[i].p_align);

			}
		}
	}

	// Parse section headers

	if(res.e_shoff && res.e_shnum) {

		if(filesize < (res.e_shnum * res.e_shentsize + res.e_shoff)) {
			throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted " +
									"(file size < (e_shnum * e_shentsize + e_shoff))");
		}

		std::vector<char> secheaders = bzbfile::read_bin(filename,
														res.e_shnum * res.e_shentsize,
														res.e_shoff);

		if(secheaders.size() < res.e_shnum * res.e_shentsize) {
			throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted " +
									"(section headers are corrupted)");
		}

		res.sec_headers.resize(res.e_shnum);
		for(size_t i = 0; i < res.sec_headers.size(); i++) {

			if(res.e_ident[ELF_EI_DATA] == ELF_LITTLE_ENDIAN) {

				readint(secheaders, 0x00 + (i * res.e_shentsize), &res.sec_headers[i].sh_name);
				readint(secheaders, 0x04 + (i * res.e_shentsize), &res.sec_headers[i].sh_type);
				readint(secheaders, 0x08 + (i * res.e_shentsize), &res.sec_headers[i].sh_flags);
				readint(secheaders, 0x10 + (i * res.e_shentsize), &res.sec_headers[i].sh_addr);
				readint(secheaders, 0x18 + (i * res.e_shentsize), &res.sec_headers[i].sh_offset);
				readint(secheaders, 0x20 + (i * res.e_shentsize), &res.sec_headers[i].sh_size);
				readint(secheaders, 0x28 + (i * res.e_shentsize), &res.sec_headers[i].sh_link);
				readint(secheaders, 0x2C + (i * res.e_shentsize), &res.sec_headers[i].sh_info);
				readint(secheaders, 0x30 + (i * res.e_shentsize), &res.sec_headers[i].sh_addralign);
				readint(secheaders, 0x38 + (i * res.e_shentsize), &res.sec_headers[i].sh_entsize);

			} else {

				readint_be(secheaders, 0x00 + (i * res.e_shentsize), &res.sec_headers[i].sh_name);
				readint_be(secheaders, 0x04 + (i * res.e_shentsize), &res.sec_headers[i].sh_type);
				readint_be(secheaders, 0x08 + (i * res.e_shentsize), &res.sec_headers[i].sh_flags);
				readint_be(secheaders, 0x10 + (i * res.e_shentsize), &res.sec_headers[i].sh_addr);
				readint_be(secheaders, 0x18 + (i * res.e_shentsize), &res.sec_headers[i].sh_offset);
				readint_be(secheaders, 0x20 + (i * res.e_shentsize), &res.sec_headers[i].sh_size);
				readint_be(secheaders, 0x28 + (i * res.e_shentsize), &res.sec_headers[i].sh_link);
				readint_be(secheaders, 0x2C + (i * res.e_shentsize), &res.sec_headers[i].sh_info);
				readint_be(secheaders, 0x30 + (i * res.e_shentsize), &res.sec_headers[i].sh_addralign);
				readint_be(secheaders, 0x38 + (i * res.e_shentsize), &res.sec_headers[i].sh_entsize);

			}
		}
	}

	if(res.e_shstrndx &&
	   res.e_shnum >= res.e_shstrndx) {

		if(res.sec_headers[res.e_shstrndx].sh_type != ELF_SHT_STRTAB) {
			throw std::runtime_error("ERROR Unable to parse \'" + filename +
						"\' because the file is corrupted (e_shstrndx points to invalid section)");
		}

		if(filesize < res.sec_headers[res.e_shstrndx].sh_offset +
					  res.sec_headers[res.e_shstrndx].sh_size) {
			throw std::runtime_error("ERROR Unable to parse \'" + filename +
						"\' because the file is corrupted (.shstrtab section header has invalid data)");
		}

		std::vector<char> shstrtab = bzbfile::read_bin(filename,
														res.sec_headers[res.e_shstrndx].sh_size,
														res.sec_headers[res.e_shstrndx].sh_offset);

		for (size_t i = 0; i < res.sec_headers.size(); ++i) {

			if(res.sec_headers[i].sh_name > res.sec_headers[res.e_shstrndx].sh_size) {
				throw std::runtime_error("ERROR Unable to parse \'" + filename +
									"\' because the file is corrupted (invalid sh_name value)");
			}

			load_name_str(shstrtab, res.sec_headers[i].name, res.sec_headers[i].sh_name);
		}

	}

	for (size_t i = 0; i < res.sec_headers.size(); ++i) {
		if(res.sec_headers[i].name == ".interp") {

			if(filesize < res.sec_headers[i].sh_offset +
					  	  res.sec_headers[i].sh_size) {
				throw std::runtime_error("ERROR Unable to parse \'" + filename +
							"\' because the file is corrupted (.interp section header has invalid data)");
			}

			std::vector<char> interp = bzbfile::read_bin(filename,
														 res.sec_headers[i].sh_size,
														 res.sec_headers[i].sh_offset);

			load_name_str(interp, res.interpreter, 0);

			break;
		}
	}

	res.filename = filename;

	return res;
}
